import zipfile
import re
import jieba
import pandas as pd
import chardet
import logging
from concurrent.futures import ThreadPoolExecutor
from tkinter import Tk, filedialog
import os

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 自定义英文停用词列表
english_stopwords = set([
    'a', 'an', 'the', 'and', 'or', 'but', 'if', 'in', 'on', 'with', 'as', 'by', 'for', 'of', 'at', 'to', 'from', 'up', 'down', 'out', 'over', 'under', 'again', 'further', 'then', 'once'
])

# 加载中文停用词
def load_stopwords(file_path):
    """从文件中加载中文停用词"""
    if not os.path.exists(file_path):
        logging.error(f"停用词文件未找到: {file_path}")
        return set()
    with open(file_path, 'r', encoding='utf-8') as f:
        stopwords = set([line.strip() for line in f])
    return stopwords

# 加载中文停用词文件
stopwords_file = r"C:\Users\dell\Desktop\python大二上\stopwords.txt"  # 替换为实际路径
chinese_stopwords = load_stopwords(stopwords_file)

def select_zip_file():
    """让用户选择 zip 文件"""
    root = Tk()
    root.withdraw()  # 隐藏主窗口
    zip_path = filedialog.askopenfilename(
        title="选择压缩包文件",
        filetypes=[("ZIP 文件", "*.zip")]
    )
    return zip_path

def extract_text_from_zip(zip_path):
    """从压缩包中读取所有文本文件"""
    texts = {}
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        with ThreadPoolExecutor() as executor:
            results = list(executor.map(lambda f: process_file(f, zip_ref), zip_ref.namelist()))
        texts = dict(results)
    return texts

def process_file(file_name, zip_ref):
    """处理单个文件"""
    if not file_name.endswith('.txt'):
        return file_name, None
    try:
        with zip_ref.open(file_name) as file:
            raw_data = file.read()
            encodings_to_try = ['utf-8', 'gbk', 'latin1', 'utf-16', 'big5']
            for encoding in encodings_to_try:
                try:
                    text = raw_data.decode(encoding)
                    logging.info(f'成功解码文件: {file_name}，编码: {encoding}')
                    return file_name, text
                except UnicodeDecodeError:
                    continue
            logging.warning(f'无法解码文件: {file_name}')
            return file_name, None
    except Exception as e:
        logging.error(f'处理文件 {file_name} 时出错: {str(e)}')
        return file_name, None

def extract_chinese_and_english(text):
    """从文本中提取中文和英文内容"""
    chinese_text = ''.join(re.findall(r'[\u4e00-\u9fff]+', text))
    english_text = ' '.join(re.findall(r'[a-zA-Z]+', text))
    return chinese_text, english_text

def count_words(text, language='chinese'):
    """统计词频"""
    if language == 'chinese':
        words = [word for word in jieba.cut(text) if word not in chinese_stopwords and len(word) > 1]
    else:
        # 简单的英文分词和停用词过滤
        words = [word.lower() for word in re.findall(r'\b\w+\b', text) if word.lower() not in english_stopwords and len(word) > 1]
    word_count = {}
    for word in words:
        word_count[word] = word_count.get(word, 0) + 1
    return word_count

def save_to_csv(word_count, output_file, encoding='utf-8'):
    """将词频统计结果保存到 CSV 文件"""
    df = pd.DataFrame(list(word_count.items()), columns=['Word', 'Count'])
    df.to_csv(output_file, index=False, encoding=encoding)

def main():
    # 让用户选择 zip 文件
    zip_path = select_zip_file()
    if not zip_path:
        logging.error("未选择文件，程序退出。")
        return

    # 处理 zip 文件
    texts = extract_text_from_zip(zip_path)
    for file_name, text in texts.items():
        if text is None:
            continue
        chinese_text, english_text = extract_chinese_and_english(text)
        chinese_word_count = count_words(chinese_text, 'chinese')
        english_word_count = count_words(english_text, 'english')
        
        base_name = os.path.splitext(file_name)[0]  # 去掉文件扩展名
        save_to_csv(chinese_word_count, f'{base_name}_chinese.csv', encoding='utf-8')
        save_to_csv(english_word_count, f'{base_name}_english.csv', encoding='utf-8')

if __name__ == "__main__":
    main()